import uuid
from typing import List, Tuple, Dict

import chromadb
from chromadb import GetResult, QueryResult
from langchain.docstore.document import Document

from vec_database._vector_type import SupportedVSType
from vec_database.base import VDBService
from vec_database.impl.chromadb_vdb_result import _get_result_to_documents, _results_to_docs_and_scores
from vec_database.config import SCORE_THRESHOLD
from vec_database.embeddings_tool.embeddings_fun_adapter import EmbeddingsFunAdapter


class ChromaVDBServiceImpl(VDBService):
    client = None
    collection = None

    def vs_type(self) -> str:
        return SupportedVSType.CHROMADB

    # 初始化扩展函数
    def do_init(self) -> None:
        if self.vs_path is not None:
            result = self.url_path.split(":")
            self.client = chromadb.HttpClient(host=result[0], port=result[1])
        else:
            self.client = chromadb.PersistentClient(path=self.vs_path)

        self.collection = self.client.get_or_create_collection(self.kb_name)

    # 增-创建知识库表
    def do_create_table(self):
        self.collection = self.client.get_or_create_collection(self.kb_name)

    # 删-删除知识库表:存在价值是被子类重写
    def do_drop_table(self):
        try:
            self.client.delete_collection(self.kb_name)
        except ValueError as e:
            if not str(e) == f"Collection {self.kb_name} does not exist.":
                raise e

    # 清-从清除知识库表向量(内容)
    def do_clear_table(self):
        pass

    # 增-向知识库添加文档
    # @param docs 添加的文档列表
    # @return 文档在向量数据库的唯一id
    def do_add_doc(self,
                   docs: List[Document],
                   **kwargs,
                   ) -> List[Dict]:
        texts = [doc.page_content for doc in docs]
        metadatas = [doc.metadata for doc in docs]
        embed_func = EmbeddingsFunAdapter(self.embed_model)
        embeddings = embed_func.embed_documents(texts)
        ids = [str(uuid.uuid1()) for _ in range(len(texts))]
        doc_infos = []
        for _id, text, embedding, metadata in zip(ids, texts, embeddings, metadatas):
            self.collection.add(ids=_id, embeddings=embedding, metadatas=metadata, documents=text)
            doc_infos.append({"id": _id, "metadata": metadata})
        return doc_infos

    # 删-根据文件路径从向量库中删除
    # @param kb_file_source 文件的路径
    def do_delete_doc(self,
                      kb_file_souce: str):
        self.collection.delete(where={"source":kb_file_souce})

    # 删-根据ids从向量库中删除
    # @param ids 向量库中唯一id集合
    def do_delete_docs_by_ids(self, ids: List[str]) -> bool:
        self.collection.delete(ids)
        return True

    # 改-修改向量内容
    def do_update_doc(self):
        pass

    # 查-根据唯一id获取内容
    # @param ids 向量数据库中唯一ids
    def get_doc_by_ids(self, ids: List[str]) -> List[Document]:
        get_result: GetResult = self.collection.get(ids)
        return _get_result_to_documents(get_result)

    # 查-根据条件搜索知识库内容
    # @param query 查询内容
    # @param top_k 最大的token数量
    # @param score_threshold 匹配的距离阈值，一般取值范围在0-1之间，SCORE越小，距离越小从而相关度越高
    def do_search_doc(self,
                      query: str,
                      top_k: int,
                      score_threshold: float = SCORE_THRESHOLD,
                      ) -> List[Tuple[Document, float]]:
        embed_func = EmbeddingsFunAdapter(self.embed_model)
        embeddings = embed_func.embed_query(query)
        query_result: QueryResult = self.collection.query(query_embeddings=embeddings, n_results=top_k)
        return _results_to_docs_and_scores(query_result)
